#include <stdint.h>
#include <stdio.h>

#include "../../inc/wlc_test.h"
#include "../../inc/dilithium/dilithium.h"
#include "../../inc/dilithium/params.h"
#include "../../inc/dilithium/sign.h"

#include "../../inc/rand/randombytes.h"

#if REE_OR_TEE == 0

#include "../../inc/wlc_bench.h"

#define MLEN 100
#define ROUND 10

void test_dilithium(int level)
{
    printf("+------------------------------------------------------------------------------+\n");
    switch (level) {
        case DILITHIUM_NIST2:
            printf("+---    Test module  --->  Dilithium Level %d.\n", 2);
            break;
        case DILITHIUM_NIST3:
            printf("+---    Test module  --->  Dilithium Level %d.\n", 3);
            break;
        case DILITHIUM_NIST5:
            printf("+---    Test module  --->  Dilithium Level %d.\n", 5);
            break;
        default:
            printf("ERROR LEVEL!!!\n");
            break;
    }
    
    dilithium_init(level);

    unsigned int i, j;
    int ret;
    size_t mlen, smlen;
    uint8_t m[MLEN] = {0};
    uint8_t sm[MLEN + MAX_CRYPTO_BYTES];
    uint8_t m2[MLEN + MAX_CRYPTO_BYTES];
    uint8_t pk[MAX_CRYPTO_PUBLICKEYBYTES];
    uint8_t sk[MAX_CRYPTO_SECRETKEYBYTES];

    for(i = 0; i < NTESTS; ++i) {
        randombytes(m, MLEN);

        crypto_sign_keypair(pk, sk);
        crypto_sign(sm, &smlen, m, MLEN, sk);
        ret = crypto_sign_open(m2, &mlen, sm, smlen, pk);

        if(ret) {
            printf("Verification failed [Fail!!!!]\n");
            goto END;
        }

        if(mlen != MLEN) {
            printf("Message lengths don't match [Fail!!!!]\n");
            goto END;
        }

        for(j = 0; j < mlen; ++j) {
            if(m[j] != m2[j]) {
                printf("Messages don't match [Fail!!!!]\n");
                goto END;
            }
        }

        randombytes((uint8_t *)&j, sizeof(j));
        do {
            randombytes(m2, 1);
        } while(!m2[0]);
        sm[j % params.CRYPTO_BYTES] += m2[0];
        ret = crypto_sign_open(m2, &mlen, sm, smlen, pk);
        if(!ret) {
            printf("Trivial forgeries possible [Fail!!!!]\n");
            goto END;
        }
    }

    printf("CRYPTO_PUBLICKEYBYTES = %d\n", params.CRYPTO_PUBLICKEYBYTES);
    printf("CRYPTO_SECRETKEYBYTES = %d\n", params.CRYPTO_SECRETKEYBYTES);
    printf("CRYPTO_BYTES = %d\n", params.CRYPTO_BYTES);

    printf("Dilithium [PASS]\n");

    END:
    printf("+------------------------------------------------------------------------------+\n");
}

void bench_dilithium(int level)
{
    printf("+------------------------------------------------------------------------------+\n");
    printf("+---    Bench module  --->  Dilithium Level ");
    dilithium_init(level);
    switch (level) {
        case DILITHIUM_NIST2:
            printf("%d.\n", 2);
            break;
        case DILITHIUM_NIST3:
            printf("%d.\n", 3);
            break;
        case DILITHIUM_NIST5:
            printf("%d.\n", 5);
            break;
        default:
            printf("ERROR LEVEL!!!\n");
            break;
    }

    unsigned int i, j;
    size_t mlen, smlen;
    uint8_t m[MLEN] = {0};
    uint8_t sm[MLEN + MAX_CRYPTO_BYTES];
    uint8_t m2[MLEN + MAX_CRYPTO_BYTES];
    uint8_t pk[MAX_CRYPTO_PUBLICKEYBYTES];
    uint8_t sk[MAX_CRYPTO_SECRETKEYBYTES];

    bench_ctx keygen, sign, vrf;
    bench_ctx_reset(&keygen);
    bench_ctx_reset(&sign);
    bench_ctx_reset(&vrf);

    for(i = 0; i < NTESTS/ROUND; ++i) {
        randombytes(m, MLEN);

        bench_ctx_before(&keygen);
        for (j = 0; j < ROUND; j++)
        {
            crypto_sign_keypair(pk, sk);
        }
        bench_ctx_after(&keygen);

        bench_ctx_before(&sign);
        for (j = 0; j < ROUND; j++)
        {
            crypto_sign(sm, &smlen, m, MLEN, sk);
        }
        bench_ctx_after(&sign);

        bench_ctx_before(&vrf);
        for (j = 0; j < ROUND; j++)
        {
            crypto_sign_open(m2, &mlen, sm, smlen, pk);
        }
        bench_ctx_after(&vrf);
    }

    bench_ctx_compute(&keygen, NTESTS);
    bench_ctx_compute(&sign, NTESTS);
    bench_ctx_compute(&vrf, NTESTS);

    bench_ctx_print(&keygen, "crypto_sign_keypair");
    bench_ctx_print(&sign, "crypto_sign");
    bench_ctx_print(&vrf, "crypto_sign_open");

    printf("+------------------------------------------------------------------------------+\n");
}


#endif // REE_OR_TEE == 0